package client

import (
	"fmt"
	"github.com/xtaci/smux"
	"io"
	"net"
	"time"
	"tinp/common"
	"tinp/mylog"
)

type Client struct {
	ClientId   string
	serverAddr string
}

func NewClient(clientID string, serverAddr string) *Client {
	return &Client{
		ClientId:   clientID,
		serverAddr: serverAddr,
	}
}

func (c *Client) Run() error {
	conn, err := net.Dial("tcp", c.serverAddr)
	if err != nil {
		return err
	}
	mylog.Logger.Info("连接服务器成功!")
	// 连接成功
	defer conn.Close()
	// 发送 handshake 包（与网关建立连接）
	handshakeReq := common.HandshakeReq{ClientId: c.ClientId}
	buf, err := handshakeReq.Encode()
	if err != nil {
		return err
	}
	// 设置写超时
	_ = conn.SetWriteDeadline(time.Now().Add(time.Second * 3))
	_, err = conn.Write(buf)
	// 取消写超时
	_ = conn.SetWriteDeadline(time.Time{})
	if err != nil {
		return err
	}
	// 创建 mux session
	muxClient, err := smux.Client(conn, nil)
	if err != nil {
		return err
	}
	defer muxClient.Close()
	// 等待 mux stream
	for {
		stream, err := muxClient.AcceptStream()
		if err != nil {
			return err
		}
		go c.handleStream(stream)
	}
}

func (c *Client) handleStream(stream *smux.Stream) {
	defer stream.Close()
	pp := &common.ProxyProtocol{}

	err := pp.Decode(stream)
	if err != nil {
		mylog.Logger.Errorf("decode pp fail: %v", err)
		return
	}
	var localCoon net.Conn
	// 与本地建立连接
	localCoon, err = net.Dial("tcp", fmt.Sprintf("%s:%d", pp.TargetIP, pp.TargetPort))
	if err != nil {
		mylog.Logger.Errorf("与代理服务建立连接失败: %v", err)
		return
	}
	defer localCoon.Close()

	// 双向数据拷贝 将服务端发送过来的消息，转发给 目标内网服务端
	go func() {
		defer func() {
			localCoon.Close()
			stream.Close()
		}()
		// 将目标内网服务端发送过来的消息，转发给外网服务端
		io.Copy(localCoon, stream)
	}()
	io.Copy(stream, localCoon)
}
